import argparse
import os
import re
import uuid

import pandas as pd
import requests
from dotenv import load_dotenv
from tqdm import tqdm

load_dotenv()

parser = argparse.ArgumentParser()
parser.add_argument("file", type=str, help="excel file path")
args = parser.parse_args()


df = pd.read_excel(args.file)
dataset_id = os.getenv("DATASET_ID")
document_id = os.getenv("DOCUMENT_ID")
api_key = os.getenv("API_KEY_DB")
base_url = os.getenv("DIFY_URL")

for _, row in tqdm(df.iterrows()):
    url = f"{base_url}/v1/datasets/{dataset_id}/documents/{document_id}/segments"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Content-Type": "application/json",
    }

    tags = row["标签"]
    if not isinstance(tags, str):
        keywords = []

    else:
        tags = [line.strip() for line in tags.split("\n") if line.strip()]
        keywords = []

        for tag in tags:
            tag = re.sub(r"一级事项引导[：:]", "", tag)
            tag = re.sub(r"二级事项引导[：:]", "", tag)
            tag = re.sub(r"场景引导[：:]", "", tag)
            tag = re.sub(r"关键词[：:]", "", tag)
            keywords.extend(re.split(r"[,，]", tag))

    # Define the data to be sent
    data = {
        "segments": [
            {
                # "content": row["问题"],
                "content": f"{'-'.join(keywords)}\n{row['问题']}\n{row['答案']}",
                "answer": row["答案"],
                "keywords": keywords,
            }
        ]
    }

    cur_url = os.path.join(url, str(uuid.UUID(row["id"])))
    # Make the POST request
    response = requests.post(cur_url, headers=headers, json=data)

    if response.status_code == 404:
        data["segments"][0]["segment_id"] = str(uuid.UUID(row["id"]))
        response = requests.post(url, headers=headers, json=data)
